import os
from typing import Annotated, Union

from langchain.chat_models import init_chat_model
from langchain.chat_models.base import _ConfigurableModel
from langchain_core.language_models import BaseChatModel
from langchain_core.tools import Tool
from langgraph.constants import START, END
from langgraph.graph import add_messages, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
from typing_extensions import TypedDict

from agents.utils import ThreadSafeSqliteSaver
from utils.utils import account_config

config = account_config()

model = config.get('model')
api_key = config.get('api_key')
base_url = config.get('base_url')

def get_weather(city):
    print(f"查询天气：{city}")
    return f"{city} 今天 26°C，晴"

class State(TypedDict):
    messages: Annotated[list, add_messages]
    llm: Union[BaseChatModel, _ConfigurableModel]


class Agent:
    chat_node_name = "chat"
    tools_node_name = "tools"

    def __init__(self, checkpointer="test.sqlite", config_id="1", tools=None):
        self.tools = tools
        self.llm = init_chat_model(model, api_key=api_key, base_url=base_url)
        thread_id = f"{config_id if config_id is not None else os.getgid()}"
        self.config = {"configurable": {"thread_id": thread_id}}
        self.graph = StateGraph(State)
        self.checkpointer = ThreadSafeSqliteSaver(checkpointer)
        self.checkpointer.delete_thread(thread_id)
        self.compile()

    @staticmethod
    def chatbot_node(state: State) -> State:
        res = state['llm'].invoke(state['messages'])
        state['messages'] = [res]
        return state

    def compile(self):

        self.add_chat_node()
        self.add_tool_node(tools=self.tools if self.tools else [])

        self.graph.add_edge(START, self.chat_node_name)
        self.graph.add_conditional_edges(self.chat_node_name, tools_condition)
        self.graph.add_edge(self.tools_node_name, self.chat_node_name)
        self.graph.add_edge(self.chat_node_name, END)

        self.graph = self.graph.compile(checkpointer=self.checkpointer)

    def draw_graph_png(self):
        # 获取图片的二进制数据
        png_data = self.graph.get_graph().draw_mermaid_png()
        # 将图片数据保存到文件
        with open("graph.png", "wb") as f:
            f.write(png_data)

    def add_chat_node(self):
        self.graph.add_node(self.chat_node_name, self.chatbot_node)

    def add_tool_node(self, tools: list):
        if not tools:
            tools = [Tool(name="get_weather", func=get_weather, description="查询天气")]

        self.llm = self.llm.bind_tools(tools)
        self.graph.add_node(self.tools_node_name, ToolNode(tools=tools))

    def stream(self, message: str):
        message = {"messages": [{"role": "user", "content": message}], "llm": self.llm}
        return self.graph.stream(message, config=self.config, stream_mode="values")

    def invoke(self, message: str):
        message = {"messages": [{"role": "user", "content": message}], "llm": self.llm}
        return self.graph.invoke(message, config=self.config)



